"""
-*- coding: utf-8 -*-
@File  : generate_filelist.py
@author: ZhenyuYang
@Time  : 2023/04/06 18:35
"""
import os

dataset = ["samm"]
cross = 5
labels = ["negative", "positive", "surprise"]
flow = True
if len(dataset) == 3:
    for c in range(cross):
        with open("../data/annotations/trainlist{:02d}.txt".format(c+1), "r") as f:
            data_list = f.readlines()
        with open("../data/MEGC2019_train_split_{}_videos.txt".format(c+1), "w") as f:
            pass
        with open("../data/MEGC2019_train_split_{}_rawframes.txt".format(c+1), "w") as f:
            pass
        if flow:
            with open("../data/MEGC2019_train_split_{}_flows.txt".format(c + 1), "w") as f:
                pass
        for data in data_list:
            data = data[:-1]
            label = labels.index(data.split("/")[0]) + 1
            with open("../data/MEGC2019_train_split_{}_videos.txt".format(c+1), "a") as f:
                f.write("{} {}\n".format(data, label-1))
            frame_count = len(os.listdir("../data/flow_prepare_rawframes/{}".format(data.split(".")[0])))
            with open("../data/MEGC2019_train_split_{}_rawframes.txt".format(c+1), "a") as f:
                f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label-1))
            if flow:
                frame_count = len(os.listdir("../data/flows_raw/{}".format(data.split(".")[0])))
                with open("../data/MEGC2019_train_split_{}_flows.txt".format(c+1), "a") as f:
                    f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))

        with open("../data/annotations/testlist{:02d}.txt".format(c+1), "r") as f:
            data_list = f.readlines()
        with open("../data/MEGC2019_val_split_{}_videos.txt".format(c+1), "w") as f:
            pass
        with open("../data/MEGC2019_val_split_{}_rawframes.txt".format(c+1), "w") as f:
            pass
        if flow:
            with open("../data/MEGC2019_val_split_{}_flows.txt".format(c + 1), "w") as f:
                pass
        for data in data_list:
            data = data[:-1]
            label = labels.index(data.split("/")[0]) + 1
            with open("../data/MEGC2019_val_split_{}_videos.txt".format(c+1), "a") as f:
                f.write("{} {}\n".format(data, label-1))
            frame_count = len(os.listdir("../data/flow_prepare_rawframes/{}".format(data.split(".")[0])))
            with open("../data/MEGC2019_val_split_{}_rawframes.txt".format(c+1), "a") as f:
                f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label-1))
            if flow:
                frame_count = len(os.listdir("../data/flows_raw/{}".format(data.split(".")[0])))
                with open("../data/MEGC2019_val_split_{}_flows.txt".format(c + 1), "a") as f:
                    f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))
elif len(dataset) == 1:
    for c in range(cross):
        with open("../data/{}_annotations/trainlist{:02d}.txt".format(dataset[0], c + 1), "r") as f:
            data_list = f.readlines()
        with open("../data/{}_train_split_{}_videos.txt".format(dataset[0].upper(), c + 1), "w") as f:
            pass
        with open("../data/{}_train_split_{}_rawframes.txt".format(dataset[0].upper(), c + 1), "w") as f:
            pass
        if flow:
            with open("../data/{}_train_split_{}_flows.txt".format(dataset[0].upper(), c + 1), "w") as f:
                pass
        for data in data_list:
            data = data[:-1]
            label = labels.index(data.split("/")[0]) + 1
            with open("../data/{}_train_split_{}_videos.txt".format(dataset[0].upper(), c + 1), "a") as f:
                f.write("{} {}\n".format(data, label - 1))
            frame_count = len(os.listdir("../data/{}_rawframes/{}".format(dataset[0], data.split(".")[0])))
            with open("../data/{}_train_split_{}_rawframes.txt".format(dataset[0].upper(), c + 1), "a") as f:
                f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))
            if flow:
                frame_count = len(os.listdir("../data/{}_flows/{}".format(dataset[0], data.split(".")[0])))
                with open("../data/{}_train_split_{}_flows.txt".format(dataset[0].upper(), c + 1), "a") as f:
                    f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))

        with open("../data/{}_annotations/testlist{:02d}.txt".format(dataset[0], c + 1), "r") as f:
            data_list = f.readlines()
        with open("../data/{}_val_split_{}_videos.txt".format(dataset[0].upper(), c + 1), "w") as f:
            pass
        with open("../data/{}_val_split_{}_rawframes.txt".format(dataset[0].upper(), c + 1), "w") as f:
            pass
        if flow:
            with open("../data/{}_val_split_{}_flows.txt".format(dataset[0].upper(), c + 1), "w") as f:
                pass
        for data in data_list:
            data = data[:-1]
            label = labels.index(data.split("/")[0]) + 1
            with open("../data/{}_val_split_{}_videos.txt".format(dataset[0].upper(), c + 1), "a") as f:
                f.write("{} {}\n".format(data, label - 1))
            frame_count = len(os.listdir("../data/{}_rawframes/{}".format(dataset[0], data.split(".")[0])))
            with open("../data/{}_val_split_{}_rawframes.txt".format(dataset[0].upper(), c + 1), "a") as f:
                f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))
            if flow:
                frame_count = len(os.listdir("../data/{}_flows/{}".format(dataset[0], data.split(".")[0])))
                with open("../data/{}_val_split_{}_flows.txt".format(dataset[0].upper(), c + 1), "a") as f:
                    f.write("{} {} {}\n".format(data.split(".")[0], frame_count, label - 1))
